﻿#################################################
#												#
# Discrete Poisson Equation 					#
# Deprecated! Use the _Epetra version instead.	#
# Matt Ball (ball@brown.edu)					#
# 2008-2009										#
#												#
#################################################

from numpy import *
from scipy import sparse, linsolve
from scipy.linalg import inv
from sympy import sympify, lambdify, Symbol, gamma
import Gnuplot

# replaced with getBoundaryFunctions
def getUserB(n):
	print "Enter the row vector b (it should have", `(n-2)**2` ,"elts)",
	try:
		b = array(input(":"),float)
	except Exception, e:
		print "\t\t",e
		print "Invalid vector entered. Make sure it is 1D and has", `(n-2)**2`,"elements."
		b = getUserB(n)
	
	if (b.ndim != 1) | (b.size != (n-2)**2):
		print "Invalid vector entered. Make sure it is 1D and has", `(n-2)**2`,"elements, not",b.size,"elts."
		b = getUserB(n)
	
	return b

# 2D heat plot using gnuplot
def plotMat(A):
	n = A[0].size
	g = Gnuplot.Gnuplot()
	x = arange(n*1.)/(n-1.)
	y = arange(n*1.)/(n-1.)
	g('set pm3d map')
	g('set size square')
	#g.title('Your title here')
	g.xlabel('x')
	g.ylabel('y')
	g.splot(Gnuplot.GridData(A.T,x,y, binary=0))
	
# 3D surface plot using gnuplot
def plotMat3D(A):
	n = A[0].size
	g = Gnuplot.Gnuplot()
	x = arange(n*1.)/(n-1.)
	y = arange(n*1.)/(n-1.)
	g('set pm3d')
	#g.title('Your title here')
	g.xlabel('x')
	g.ylabel('y')
	g('set size square')
	g('unset surface')
	g.splot(Gnuplot.GridData(A.T,x,y, binary=0))
	
def getPoissonF():
	print "Enter f in ∇²u = f (0 for Laplace case). f is a 2D equation in terms of x and y."
	f = sympify(raw_input("f = "))
	symbols = f.atoms(Symbol)
	
	x = Symbol("x")
	y = Symbol("y")
	
	# Make sure x and/or y is actually a symbol in the equation
	if (not len(symbols) == 0):
		if x in symbols:
			if y in symbols:
				if len(symbols) == 2:
					return lambdify((x,y),f.subs({'x':x,'y':y}))
				else:
					print "\tYou entered an equation with too many variables, try again."
					return getPoissonF()
			else:
				if len(symbols) == 1:
					return lambdify((x,y),f.subs({'x':x}))
				else:
					print "\tTry again."
					return getPoissonF()
		else:
			if y in symbols:
				if len(symbols) == 1:
					return lambdify((x,y),f.subs({'y':y}))
				else:
					print "\tTry again."
					return getPoissonF()
			else:
				return lambdify((x,y),f.subs({'x':x,'y':y}))
	else:
		return lambdify((x,y),f)
	
def getBoundaryFunctions():
	print "Boundary conditions: enter 1D equations where x is the independent variable."
	x = Symbol("x")
				
	def getFunctionWithX(message):
		try:
			g = sympify(raw_input(message))
			symbols = g.atoms(Symbol)
	
			# Make sure x is actually a symbol in the equation
			if ((x in symbols) and (len(symbols)==1)) or len(symbols)==0:
				return g
			else:
				print "\tx is not a variable in the equation you entered, or you entered a higher-dimensional equation."
				return getFunctionWithX(message)
		except Exception,e:
			print "\t",e
			print "\ttry again:"
			return getFunctionWithX(message)
	
	g_top = sympify(getFunctionWithX("Enter boundary equation for the top: g_top = "))
	g_bottom = sympify(getFunctionWithX("Enter boundary equation for the bottom: g_bottom = "))
	g_left = sympify(getFunctionWithX("Enter boundary equation for the left: g_left = "))
	g_right = sympify(getFunctionWithX("Enter boundary equation for the right: g_right = "))
	
	# Convert to python lambda	
	g_top = lambdify(x,g_top.subs({'x':x}))
	g_bottom = lambdify(x,g_bottom.subs({'x':x}))
	g_left = lambdify(x,g_left.subs({'x':x}))
	g_right = lambdify(x,g_right.subs({'x':x}))
	
	# Now they can be used to actually compute stuff
	functions = [g_top, g_bottom, g_left, g_right]
	return functions

def checkContinuity(functions):
	g_top = functions[0]
	g_bottom = functions[1]
	g_left = functions[2]
	g_right = functions[3]
	
	# top left (0,1)
	if (g_top(0) == g_left(1)):
		topLeft = True
	else:
		topLeft = False
		print "\tBoundary functions are not continuous at top left corner."
	
	# top right (1,1)
	if (g_top(1) == g_right(1)):
		topRight = True
	else:
		topRight = False
		print "\tBoundary functions are not continuous at top right corner."
	
	# bottom left (0,0)
	if (g_bottom(0) == g_left(0)):
		bottomLeft = True
	else:
		bottomLeft = False
		print "\tBoundary functions are not continuous at bottom left corner."
	
	# bottom right (1,0)
	if (g_bottom(1) == g_right(0)):
		bottomRight = True
	else:
		bottomRight = False
		print "\tBoundary functions are not continuous at bottom right corner."
	
	return (topLeft and topRight and bottomLeft and bottomRight)
	
# RMS error, more or less
def calcError(A,poissonF):
	# (1/(n-2)²)*[sum over i,j (average of surrounding points - u_i,j - f_i,j)²]
	sum = 0
	n = A[0].size
	
	for i in range(1,n-1):
		for j in range(1,n-1):
			N = A[i,j+1]
			S = A[i,j-1]
			E = A[i+1,j]
			W = A[i-1,j]
			f = poissonF(i/n,j/n)
			
			sum = sum + ((N+S+E+W)/4 - A[i,j] - f)**2

	return sum/((n-2)**2)

def main():
	
	n = int(raw_input("Enter an int value for n: "))
	
	if n <= 4:
		print "\tThis value of n is not recommended, proceeding anyway. Your results may vary."
	
	# pretty useful constant
	blockSize = n-2
	
	###########
	# Make A: #
	###########
	A = zeros((blockSize**2,blockSize**2))
	
	# set the diagonals
	for i in range(blockSize**2):
		A[i,i] = 4
	
	# set the off-diagonals
	for i in range(blockSize**2 - blockSize):
		A[i,i+blockSize] = -1
		A[i+blockSize,i] = -1

	# set the blocks around the diagonal
	for i in range(blockSize**2):
		mod = i % blockSize
		if ((mod>0) and (mod<blockSize-1)):
			try:
				A[i-1,i] = -1
			except IndexError:
				pass
			try:	
				A[i+1,i] = -1
			except IndexError:
				pass
			try:	
				A[i,i-1] = -1
			except IndexError:
				pass
			try:	
				A[i,i+1] = -1
			except IndexError:
				pass

	###########
	# Make b: #
	###########
	b = zeros(blockSize**2)
	h = 1./(n-1) # used in boundary functions
	
	# get boundary stuff
	poissonF = getPoissonF()
	functions = getBoundaryFunctions()
	
	# check continuity @ corners
	while (not checkContinuity(functions)):
		print "\tRe-enter boundary condition functions."
		functions = getBoundaryFunctions()

	g_top = functions[0]
	g_bottom = functions[1]
	g_left = functions[2]
	g_right = functions[3]
	
	# helper function for mapping u_(i,j) in the "grid" to elements of x (or b)
	def xIndex(i,j):
		return (i-1) + (j-1)*blockSize
	
	# fill b
	for i in range(1,blockSize+1):
		for j in range(1,blockSize+1):
			b[xIndex(i,j)] = b[xIndex(i,j)] - (h**2)*poissonF(i*h,j*h)
			if (i==1):
				b[xIndex(i,j)] = b[xIndex(i,j)] + g_left(j*h)
			if (j==1):
				b[xIndex(i,j)] = b[xIndex(i,j)] + g_bottom(i*h)
			if (i==blockSize):
				b[xIndex(i,j)] = b[xIndex(i,j)] + g_right(j*h)
			if (j==blockSize):
				b[xIndex(i,j)] = b[xIndex(i,j)] + g_top(i*h)
	
	result_ = dot(inv(A),b)
	result = result_.reshape(-1,blockSize)
	result_ = dot(A,result_)
	
#	print "A = \n", A
# 	print "b = \n", b
#	print "result = \n", result
#	print "Ax = \n", result_

	# construct the view-friendly result matrix
	toView = zeros((n,n))
 	
 	# put in boundary functions
	for i in range(n):
		toView[n-1,i] = g_top(i*h)
		toView[0,i] = g_bottom(i*h)
		toView[i,0] = g_left(i*h)
		toView[i,n-1] = g_right(i*h)
		
	# embed result in toView
 	for i in range(blockSize):
 		for j in range(blockSize):
 			toView[i+1,j+1] = result[i,j]
 	
 	print "error =", `calcError(toView,poissonF)`
 
 	plotMat3D(toView)
 	plotMat(toView)

main()